# *-* coding:utf8 *-*
import frida
import sys
from pathlib import Path
import time


def analysis_method_name(text):
    text = text.split('(')
    text = text[0]
    start = text.rfind('.')
    return text[start + 1:]


def on_message(message, data):
    if message['type'] == 'send':
        text = message['payload']
        print(text)
    else:
        print(message)


class FridaManager:
    def __init__(self, log_out=None, is_usb=True, host=''):
        self._log_out = log_out
        self._is_usb = is_usb
        self._host = host
        self._script = None
        self._script_buffer = None
        self._device = None
        self._session = None
        self._pid = 0

    def init_device(self):
        if self._is_usb:
            try:
                self._device = frida.get_usb_device()
            except Exception as e:
                self._device = frida.get_remote_device()
        else:
            if self._host:
                manager = frida.get_device_manager()
                self._device = manager.add_remote_device(self._host)
            else:
                self._device = frida.get_local_device()

    def set_devices(self):
        pass

    def attach_app(self, process):
        self._session = self._device.attach(process)

    def attach_load(self, callback=on_message):
        self._script = self._session.create_script(self._script_buffer)
        self._script.on('message', callback)
        self._script.load()
        # sys.stdin.read()

    def swap_app(self, process):
        self._pid = self._device.spawn(process)
        self._session = self._device.attach(self._pid)

    def swap_load(self, callback=on_message):
        self._script = self._session.create_script(self._script_buffer)
        self._script.on('message', callback)
        self._script.load()
        self._device.resume(self._pid)
        time.sleep(2)
        sys.stdin.read()

    def unload(self):
        self._script.unload()
        self._script = None

    def detach(self):
        self._session.detach()
        self._script = None

    def get_exports(self):
        if self._script:
            return self._script.exports
        else:
            return None

    def load_script(self, script_path):
        with open(script_path, encoding="utf-8") as f:
            self._script_buffer = f.read()

    def get_applications(self):
        try:
            if not self._device:
                return None
            result = []
            applications = self._device.enumerate_applications()
            if len(applications) > 0:
                for application in sorted(applications, key=cmp_to_key(compare_applications)):
                    # if application.pid != 0:
                    # print(application.pid, application.name, application.identifier)
                    result.append(application.name)
        except Exception as e:
            print(e)
        return result

    def get_processes(self):
        try:
            if not self._device:
                return None
            result = []
            processes = self._device.enumerate_processes()
            if len(processes) > 0:
                for process in sorted(processes, key=cmp_to_key(compare_applications)):
                    result.append(str(process.pid) + ' ' + process.name)
        except Exception as e:
            print(e)
        return result


def compare_applications(a, b):
    a_is_running = a.pid != 0
    b_is_running = b.pid != 0
    if a_is_running == b_is_running:
        if a.name > b.name:
            return 1
        elif a.name < b.name:
            return -1
        else:
            return 0
    elif a_is_running:
        return -1
    else:
        return 1


def cmp_to_key(cmp):
    """Convert a cmp= function into a key= function"""

    class K:
        def __init__(self, obj):
            self.obj = obj

        def __lt__(self, other):
            return cmp(self.obj, other.obj) < 0

        def __gt__(self, other):
            return cmp(self.obj, other.obj) > 0

        def __eq__(self, other):
            return cmp(self.obj, other.obj) == 0

        def __le__(self, other):
            return cmp(self.obj, other.obj) <= 0

        def __ge__(self, other):
            return cmp(self.obj, other.obj) >= 0

        def __ne__(self, other):
            return cmp(self.obj, other.obj) != 0

    return K


if __name__ == '__main__':
    frida_tools = FridaManager()
    frida_tools.init_device()
    frida_tools.load_script('C:/Users/Administrator/Desktop/frida-script-project/_agent.js')
    frida_tools.swap_app('com.tencent.mobileqq')#com.tencent.mobileqq
    frida_tools.swap_load(on_message)
    # frida_tools.get_processes()